from PIL import Image
from tools_zsj.ZsjDataSet import SingleImageDataSet
import pickle
import numpy as np
import threading

# 从SegmentationClass目录下的图片中获取到各种颜色对应的目标类型, 如黑色应该代表背景，白色代表物体边界，红色是飞机等等

color_object_dict = dict()
now_id = 0
progress = 1
threadingLock = threading.Lock()

dataset = SingleImageDataSet()
dataset.push_image_folder("/home/xiaoxiang/DataSet/VOC2012/SegmentationClass") #这里设置PASCAL VOC数据集中SegmentationClass文件夹的路径
threads = []


def collect_from_one_image():
    global progress
    global dataset
    global now_id
    global color_object_dict

    while True:
        print("准备第", progress, "个...")
        threadingLock.acquire()
        own = progress
        progress = progress+1
        threadingLock.release()

        image_data = dataset.read_image_data(dataset.data[own])
        width = image_data.width
        height = image_data.height
        image_data = np.array(image_data)
        for row in range(0, height):
            for col in range(0, width):
                label = image_data[row][col].tostring()
                if label in color_object_dict.keys():
                    continue
                else:
                    threadingLock.acquire()
                    print(str(image_data[row][col]) + " : " + str(now_id))
                    color_object_dict[label] = now_id
                    now_id = now_id + 1
                    threadingLock.release()
        if now_id >= 22:
            break
        print("第", own, "个已完成")


if __name__ == '__main__':
    for index in range(1): # 使用15个线程同时运行，提高效率
        try:
            th = threading.Thread(target = collect_from_one_image)
            threads.append(th)
            th.start()
        except:
            print("线程启动失败")
    for t in threads:
        t.join()


    with open("color_object_dict.pickle", 'wb') as f: # 保存获取到的颜色值与对应类别的字典
        pickle.dump(color_object_dict, f)